#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import matplotlib.pyplot as plt
import rospy
from manip_2l.msg import EndEffectorPos, JointAngles
from model.manipulator import Manipulator


class ManipulatorVisualizer:
    def __init__(self):
        # Initialize the figure
        fig = plt.figure()
        self.ax = fig.add_subplot()

    def draw_robot(self, manipulator, theta1, theta2):
        '''
        Draw the robot given the robot model and joint values

        manipulator:  The robot instance (class Manipulator)
        theta1: The value of joint1
        theta2: The value of joint2
        '''

        p_l1 = manipulator.get_link1_pos(theta1, theta2)
        p_ee = manipulator.get_ee_pos(theta1, theta2)

        # Clear the current axes
        self.ax.clear()

        # Show the robot angles on the title
        text = r'$\theta_1$={0:.3f} | $\theta_2$={1:.3f}'.format(theta1, theta2)
        self.ax.set_title(text)

        # Draw the robot
        self.ax.plot([0, p_l1[0]], [0, p_l1[1]], color='blue', linewidth=4, zorder=1)
        self.ax.plot([p_l1[0], p_ee[0]], [p_l1[1], p_ee[1]], color='lightblue', linewidth=4, zorder=1)
        # Draw ground
        self.ax.plot([-5, 5], [0, 0], color='gray', linewidth=1, zorder=1)
        # Draw base
        self.ax.plot([-0.1, 0, 0.1, -0.1], [0, 0.173, 0, 0], color='gray', linewidth=4, zorder=2)
        # Draw points
        j1 = plt.Circle((p_l1[0], p_l1[1]), 0.16, fill=True, zorder=3)
        j2 = plt.Circle((p_ee[0], p_ee[1]), 0.12, fill=True, zorder=3)
        self.ax.add_artist(j1)
        self.ax.add_artist(j2)

        # Set the limit to x-axis and y-axis
        self.ax.set_xlim(-4, 4)
        self.ax.set_ylim(-2, 6)
        self.ax.set_aspect('equal', adjustable='box')

        plt.draw()
        plt.pause(0.001)

class VisualizerApp:
    def __init__(self, sub_topic='joint_angles'):
        self.__sub_topic = sub_topic
        self.manipulator = Manipulator()
        self.visualizer = ManipulatorVisualizer()
        self.__ee_pub = rospy.Publisher('/manipulator/ee_pos', EndEffectorPos, queue_size=1)

        rospy.Subscriber(self.__sub_topic, JointAngles, self.__joint_angles_callback, queue_size=1)

    def __joint_angles_callback(self, msg):
        eep = EndEffectorPos()
        eep.header = msg.header

        self.manipulator.t1 = msg.joint_angles[0]
        self.manipulator.t2 = msg.joint_angles[1]
        xy = self.manipulator.get_ee_pos(self.manipulator.t1, self.manipulator.t2)

        eep.x = xy[0]
        eep.y = xy[1]

        rospy.loginfo('{}, {}'.format(self.manipulator.t1, self.manipulator.t2))

        self.__ee_pub.publish(eep)


if __name__ == '__main__':
    rospy.init_node('manipulator_visualizer', anonymous=False)

    plt.ion()
    plt.show()

    rospy.loginfo('Manipulator visualizer loaded.')
    vis_app = VisualizerApp()

    while not rospy.is_shutdown():
        vis_app.visualizer.draw_robot(vis_app.manipulator, vis_app.manipulator.t1, vis_app.manipulator.t2)

    rospy.loginfo('END Manipulator visualizer.')








